
import os, random
import PIL.Image as Image
import multiprocessing

from utils import universal_util
from configs import changeFolders_config


def run_proc(name):
    # 数据集比例
    rate_train = 0.6
    rate_val = 0.2
    rate_test = 0.2

    picsPath = changeFolders_config.FromPath + name + "/"
    picNames = os.listdir(picsPath)
    random.shuffle(picNames)  # 打乱图像

    # 计算train\val\test数量
    length = len(picNames)
    num_train = int(rate_train * length)
    num_val = int(rate_val * length)
    num_test = length - num_train - num_val

    # 转存train部分
    for j in range(0, num_train):
        rawPath = picsPath + picNames[j]
        newPath = changeFolders_config.ToPath + "train/" + name + "/"
        universal_util.MakeFloder(newPath)
        image = Image.open(rawPath)  # 读取
        newNames = newPath + picNames[j]
        image.save(newNames)  # 转存

    # 转存val部分
    for j in range(num_train, num_train+num_test):
        rawPath = picsPath + picNames[j]
        newPath = changeFolders_config.ToPath + "val/" + name + "/"
        universal_util.MakeFloder(newPath)
        image = Image.open(rawPath)  # 读取
        newNames = newPath + picNames[j]
        image.save(newNames)  # 转存

    # 转存test部分
    for j in range(num_train+num_test, length):
        rawPath = picsPath + picNames[j]
        newPath = changeFolders_config.ToPath + "test/" + name + "/"
        universal_util.MakeFloder(newPath)
        image = Image.open(rawPath)  # 读取
        newNames = newPath + picNames[j]
        image.save(newNames)  # 转存

'''
    转存
'''
def changeFloders():
    specialNames = os.listdir(changeFolders_config.FromPath)

    pool = multiprocessing.Pool(processes=16)
    pool.map_async(run_proc, specialNames)
    pool.close()
    pool.join()





